import streamlit as st
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import io

# 设置页面配置
st.set_page_config(page_title="COVID-19 数据分析", page_icon="📊", layout="wide")

# 页面标题和描述
st.title('COVID-19 数据分析工具')
st.write('上传并分析COVID-19疫情数据')

# 数据上传
uploaded_file = st.file_uploader("上传CSV文件", type=['csv'])
if uploaded_file is not None:
    # 读取数据
    try:
        data = pd.read_csv(uploaded_file)
        st.write(data.head())  # 显示数据预览
        st.write("列名:", data.columns.tolist())  # 打印列名
    except Exception as e:
        st.error(f'数据加载失败: {e}')

    # 数据预处理
    if 'dateId' in data.columns:
        data['dateId'] = pd.to_datetime(data['dateId'], format='%Y%m%d')
        data = data.sort_values(by='dateId')
        st.write(data.head())  # 显示清洗后的数据预览

    # 数据分析和可视化
    if not data.empty:
        # 总感染趋势图
        fig1 = plt.figure(figsize=(10, 6))
        plt.plot(data['dateId'], data['confirmedCount'], marker='o', label='累计确诊')
        plt.title('总感染趋势图')
        plt.xlabel('日期')
        plt.ylabel('累计确诊人数')
        plt.xticks(rotation=45)
        plt.legend()
        st.pyplot(fig1)

        # 各城市感染曲线图
        if 'countryName' in data.columns:
            fig2 = plt.figure(figsize=(10, 6))
            for country in data['countryName'].unique():
                country_data = data[data['countryName'] == country]
                plt.plot(country_data['dateId'], country_data['confirmedCount'], marker='o', label=country)
            plt.title('各城市感染人数随时间变化曲线图')
            plt.xlabel('日期')
            plt.ylabel('感染人数')
            plt.xticks(rotation=45)
            plt.legend()
            st.pyplot(fig2)

        # 风险区域识别
        threshold = st.slider('设置感染人数阈值', min_value=100, max_value=10000, value=1000, step=100)
        if 'countryName' in data.columns:
            high_risk = data[data['confirmedCount'] > threshold]
            fig3 = plt.figure(figsize=(10, 6))
            for country in high_risk['countryName'].unique():
                country_data = high_risk[high_risk['countryName'] == country]
                plt.plot(country_data['dateId'], country_data['confirmedCount'], marker='o', label=country, color='red')
            plt.title('高风险区域感染人数随时间变化曲线图')
            plt.xlabel('日期')
            plt.ylabel('感染人数')
            plt.xticks(rotation=45)
            plt.legend()
            st.pyplot(fig3)

        # 结果输出
        # 保存图表为图片并提供下载链接
        buf = io.BytesIO()
        fig1.savefig(buf, format="png")
        buf.seek(0)
        st.download_button(
            label="下载总感染趋势图",
            data=buf,
            file_name="total_trend.png",
            mime="image/png",
        )

# 运行Streamlit应用
if __name__ == '__main__':
    st.write("欢迎使用COVID-19数据分析工具")